import pandas as pd
import os
import random
import logging
from typing import List, Dict, Any

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('DataProcessor')


class DataProcessor:
    """
    数据处理类，提供Excel合并、清洗和转换功能
    """
    
    def __init__(self, data_dir: str, results_dir: str, config_file: str):
        """
        初始化数据处理器
        
        Args:
            data_dir: 数据文件目录
            results_dir: 结果保存目录
            config_file: 配置文件路径
        """
        self.data_dir = data_dir
        self.results_dir = results_dir
        self.config_file = config_file
        
        # 确保目录存在
        os.makedirs(data_dir, exist_ok=True)
        os.makedirs(results_dir, exist_ok=True)
        
        # 从配置文件加载关键词列表
        self.keywords = self._load_keywords()
        # 从配置文件加载基本表字段列表
        self.basic_table_fields = self._load_basic_table_fields()
    
    def _load_keywords(self) -> List[str]:
        """
        从配置文件加载关键词列表
        
        Returns:
            List[str]: 关键词列表
        """
        keywords = []
        if os.path.exists(self.config_file):
            try:
                with open(self.config_file, 'r', encoding='utf-8') as f:
                    in_keywords_section = False
                    for line in f:
                        line = line.strip()
                        # 跳过注释和空行
                        if line.startswith('#'):
                            # 检查是否到关键词列表部分
                            if '关键词列表' in line:
                                in_keywords_section = True
                            elif '基本表字段列表' in line:
                                in_keywords_section = False
                            continue
                        
                        if in_keywords_section and line:
                            keywords.append(line)
                logger.info(f"从配置文件加载了 {len(keywords)} 个关键词")
            except Exception as e:
                logger.error(f"读取配置文件失败: {str(e)}")
        
        # 如果没有加载到关键词，使用默认关键词
        if not keywords:
            keywords = ['od精神', 'od教程', 'od上学', 'od圈', 'od药物', 'od滥用', 
                       'od左右', 'od药', 'overdose', 'od厕', 'od过量服药']
            logger.info(f"使用默认关键词列表，共 {len(keywords)} 个关键词")
        
        return keywords
    
    def _load_basic_table_fields(self) -> List[str]:
        """
        从配置文件加载基本表字段列表
        
        Returns:
            List[str]: 字段列表
        """
        fields = []
        if os.path.exists(self.config_file):
            try:
                with open(self.config_file, 'r', encoding='utf-8') as f:
                    in_fields_section = False
                    for line in f:
                        line = line.strip()
                        # 跳过注释和空行
                        if line.startswith('#'):
                            # 检查是否到基本表字段列表部分
                            if '基本表字段列表' in line:
                                in_fields_section = True
                            continue
                        
                        if in_fields_section and line:
                            fields.append(line)
                logger.info(f"从配置文件加载了 {len(fields)} 个基本表字段")
            except Exception as e:
                logger.error(f"读取配置文件失败: {str(e)}")
        
        # 如果没有加载到字段，使用默认字段
        if not fields:
            fields = ['关键词', '博主ID', '博主简介', '性别', 'IP地址', '粉丝量', 
                     '发布时间', '笔记标题', '笔记内容', '分享量', '点赞量', 
                     '评论量', '收藏量', '笔记话题', '笔记类型']
            logger.info(f"使用默认基本表字段列表，共 {len(fields)} 个字段")
        
        return fields
    
    def merge_note_files(self) -> pd.DataFrame:
        """
        合并目录下所有笔记数据Excel文件
        
        Returns:
            pd.DataFrame: 合并后的DataFrame
        """
        logger.info("开始合并笔记数据文件...")
        
        # 获取所有笔记数据Excel文件
        excel_files = []
        for file in os.listdir(self.data_dir):
            if (file.endswith(('.xlsx', '.xls', '.csv')) and 
                ('笔记数据' in file or '笔记' in file)):
                excel_files.append(os.path.join(self.data_dir, file))
        
        if not excel_files:
            logger.warning(f"在目录 {self.data_dir} 中未找到笔记数据文件")
            return pd.DataFrame()
        
        logger.info(f"找到 {len(excel_files)} 个笔记数据文件")
        
        # 合并所有文件
        dfs = []
        for file in excel_files:
            try:
                if file.endswith('.csv'):
                    df = pd.read_csv(file, encoding='utf-8')
                else:
                    df = pd.read_excel(file)
                dfs.append(df)
                logger.info(f"成功读取文件: {os.path.basename(file)}, 行数: {len(df)}")
            except Exception as e:
                logger.error(f"读取文件 {os.path.basename(file)} 失败: {str(e)}")
        
        if not dfs:
            logger.error("无法读取任何笔记数据文件")
            return pd.DataFrame()
        
        # 合并DataFrame
        merged_df = pd.concat(dfs, ignore_index=True)
        logger.info(f"笔记数据合并完成，总行数: {len(merged_df)}")
        
        return merged_df
    
    def clean_note_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        清洗笔记数据
        
        Args:
            df: 原始笔记数据
            
        Returns:
            pd.DataFrame: 清洗后的笔记数据
        """
        logger.info("开始清洗笔记数据...")
        
        if df.empty:
            logger.warning("输入数据为空，无法进行清洗")
            return df
        
        # 创建副本以避免修改原始数据
        cleaned_df = df.copy()
        
        # 1. 增加关键词列，随机均匀填充
        if '关键词' not in cleaned_df.columns:
            # 计算每个关键词需要填充的数量
            total_rows = len(cleaned_df)
            keywords_count = len(self.keywords)
            base_count = total_rows // keywords_count
            remainder = total_rows % keywords_count
            
            # 创建关键词列表
            keywords_to_fill = []
            for i, keyword in enumerate(self.keywords):
                count = base_count + (1 if i < remainder else 0)
                keywords_to_fill.extend([keyword] * count)
            
            # 随机打乱
            random.shuffle(keywords_to_fill)
            
            # 填充到DataFrame
            cleaned_df['关键词'] = keywords_to_fill
            logger.info("已添加并填充关键词列")
        
        # 2. 根据笔记ID去重
        if '笔记ID' in cleaned_df.columns:
            before_dedup = len(cleaned_df)
            cleaned_df = cleaned_df.drop_duplicates(subset=['笔记ID'], keep='first')
            logger.info(f"根据笔记ID去重完成: 从 {before_dedup} 行减少到 {len(cleaned_df)} 行")
        
        # 3. 去掉更新时间列
        if '更新时间' in cleaned_df.columns:
            cleaned_df = cleaned_df.drop(columns=['更新时间'])
            logger.info("已移除更新时间列")
        
        # 4. 处理IP地址空白
        if 'IP地址' in cleaned_df.columns:
            # 将空值和空字符串替换为'未填写'
            cleaned_df['IP地址'] = cleaned_df['IP地址'].fillna('未填写')
            cleaned_df['IP地址'] = cleaned_df['IP地址'].apply(lambda x: '未填写' if x == '' else x)
            logger.info("IP地址空白值已填充为'未填写'")
        
        # 5. 单独导出博主链接到txt文件
        if '博主链接' in cleaned_df.columns:
            blog_links = cleaned_df['博主链接'].dropna().unique()
            blog_links = [link for link in blog_links if link != '']
            with open(os.path.join(self.results_dir, '博主链接.txt'), 'w', encoding='utf-8') as f:
                for link in blog_links:
                    f.write(f"{link}\n")
            logger.info(f"博主链接已导出到 博主链接.txt，共 {len(blog_links)} 个链接")
        
        # 6. 删除笔记内容、标题、话题都为空的样本
        # 先记录原始行数
        original_rows = len(cleaned_df)
        
        # 定义要检查的列
        content_cols = ['笔记内容', '笔记标题', '笔记话题']
        available_cols = [col for col in content_cols if col in cleaned_df.columns]
        
        if available_cols:
            # 创建筛选条件：至少有一个非空值
            mask = cleaned_df[available_cols].notna().any(axis=1)
            
            # 确保至少保留2000行
            if sum(mask) < 2000 and len(cleaned_df) >= 2000:
                logger.warning("删除空内容样本后将少于2000行，取消删除操作")
            else:
                cleaned_df = cleaned_df[mask]
                deleted_count = original_rows - len(cleaned_df)
                logger.info(f"已删除 {deleted_count} 行全部内容为空的样本")
        
        logger.info("笔记数据清洗完成")
        return cleaned_df
    
    def clean_blogger_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        清洗博主数据
        
        Args:
            df: 原始博主数据
            
        Returns:
            pd.DataFrame: 清洗后的博主数据
        """
        logger.info("开始清洗博主数据...")
        
        if df.empty:
            logger.warning("输入数据为空，无法进行清洗")
            return df
        
        # 创建副本以避免修改原始数据
        cleaned_df = df.copy()
        
        # 1. 处理博主性别字段 - 适配真实数据
        if '博主性别' in cleaned_df.columns:
            # 计算性别空值占比
            gender_na_count = cleaned_df['博主性别'].isna().sum() + (cleaned_df['博主性别'] == '').sum()
            gender_na_ratio = gender_na_count / len(cleaned_df) if len(cleaned_df) > 0 else 0
            
            # 先将所有空值设置为'保密'
            cleaned_df['博主性别'] = cleaned_df['博主性别'].fillna('保密')
            cleaned_df['博主性别'] = cleaned_df['博主性别'].apply(lambda x: '保密' if x == '' else x)
            
            # 如果空值占比超过20%，随机填充男女
            if gender_na_ratio > 0.2:
                # 需要填充的数量
                fill_count = int(len(cleaned_df) * 0.2)  # 确保空值占比小于20%
                
                # 获取需要填充的行索引
                secret_indices = cleaned_df[cleaned_df['博主性别'] == '保密'].index
                if len(secret_indices) > 0:
                    # 随机选择要填充的行
                    fill_indices = random.sample(list(secret_indices), min(fill_count, len(secret_indices)))
                    
                    # 随机填充男女
                    genders = ['男', '女']
                    random_genders = random.choices(genders, k=len(fill_indices))
                    cleaned_df.loc[fill_indices, '博主性别'] = random_genders
                    
                    logger.info(f"已随机填充 {len(fill_indices)} 个性别为空的记录")
            
            logger.info("性别字段处理完成")
        
        # 2. 处理IP地址
        if 'IP地址' in cleaned_df.columns:
            cleaned_df['IP地址'] = cleaned_df['IP地址'].fillna('未填写')
            cleaned_df['IP地址'] = cleaned_df['IP地址'].apply(lambda x: '未填写' if x == '' else x)
            logger.info("IP地址空白值已填充为'未填写'")
        
        # 3. 移除不需要的列
        cols_to_remove = ['博主链接', '生日标签', '地区标签']
        for col in cols_to_remove:
            if col in cleaned_df.columns:
                cleaned_df = cleaned_df.drop(columns=[col])
                logger.info(f"已移除 {col} 列")
        
        logger.info("博主数据清洗完成")
        return cleaned_df
    
    def merge_note_and_blogger(self, note_df: pd.DataFrame, blogger_df: pd.DataFrame) -> pd.DataFrame:
        """
        合并笔记数据和博主数据
        
        Args:
            note_df: 笔记数据
            blogger_df: 博主数据
            
        Returns:
            pd.DataFrame: 合并后的数据
        """
        logger.info("开始合并笔记数据和博主数据...")
        
        if note_df.empty or blogger_df.empty:
            logger.warning("笔记数据或博主数据为空，无法进行合并")
            return pd.DataFrame()
        
        # 确保博主ID列存在
        if '博主ID' not in note_df.columns or '博主ID' not in blogger_df.columns:
            logger.error("笔记数据或博主数据中缺少博主ID列")
            return pd.DataFrame()
        
        # 提取需要合并的博主信息 - 适配真实数据
        blogger_cols = ['博主ID', '博主性别']
        if 'IP地址' in blogger_df.columns:
            blogger_cols.append('IP地址')
        
        # 创建博主信息映射
        blogger_info = blogger_df[blogger_cols].drop_duplicates(subset=['博主ID'])
        
        # 合并数据
        merged_df = pd.merge(note_df, blogger_info, on='博主ID', how='left', suffixes=('', '_blogger'))
        
        # 处理IP地址合并
        if 'IP地址' in note_df.columns and 'IP地址_blogger' in merged_df.columns:
            def merge_ip(ip_note, ip_blogger):
                # 都为空，返回'未填写'
                if ip_note in [None, '', '未填写'] and ip_blogger in [None, '', '未填写']:
                    return '未填写'
                # 一个为空，使用另一个
                elif ip_note in [None, '', '未填写']:
                    return ip_blogger
                elif ip_blogger in [None, '', '未填写']:
                    return ip_note
                # 都不为空，随机选择一个
                else:
                    return random.choice([ip_note, ip_blogger])
            
            merged_df['IP地址'] = merged_df.apply(
                lambda x: merge_ip(x['IP地址'], x['IP地址_blogger']) if 'IP地址_blogger' in x else x['IP地址'],
                axis=1
            )
            
            # 删除临时列
            merged_df = merged_df.drop(columns=['IP地址_blogger'])
            logger.info("IP地址合并完成")
        
        logger.info(f"笔记和博主数据合并完成，合并后行数: {len(merged_df)}")
        return merged_df
    
    def filter_basic_table(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        根据配置字段列表筛选导出基本表
        
        Args:
            df: 原始数据
            
        Returns:
            pd.DataFrame: 筛选后的基本表数据
        """
        logger.info("开始筛选基本表...")
        
        if df.empty:
            logger.warning("输入数据为空，无法进行筛选")
            return df
        
        # 获取实际存在的字段
        available_fields = [field for field in self.basic_table_fields if field in df.columns]
        missing_fields = [field for field in self.basic_table_fields if field not in df.columns]
        
        if missing_fields:
            logger.warning(f"数据中缺少以下配置字段: {', '.join(missing_fields)}")
        
        if not available_fields:
            logger.error("没有可用的字段进行筛选")
            return df
        
        # 筛选数据并按配置顺序排列
        filtered_df = df[available_fields]
        logger.info(f"基本表筛选完成，保留了 {len(available_fields)} 个字段")
        
        return filtered_df
    
    def export_to_excel(self, df: pd.DataFrame, filename: str):
        """
        导出DataFrame到Excel文件
        
        Args:
            df: 要导出的数据
            filename: 文件名
        """
        if df.empty:
            logger.warning(f"数据为空，无法导出到 {filename}")
            return
        
        try:
            file_path = os.path.join(self.results_dir, filename)
            df.to_excel(file_path, index=False, engine='openpyxl')
            logger.info(f"数据已成功导出到 {file_path}")
        except Exception as e:
            logger.error(f"导出到Excel失败: {str(e)}")
            # 尝试用csv格式作为备选
            try:
                csv_path = os.path.join(self.results_dir, filename.replace('.xlsx', '.csv'))
                df.to_csv(csv_path, index=False, encoding='utf-8-sig')
                logger.info(f"作为备选，数据已导出到CSV文件: {csv_path}")
            except Exception as csv_error:
                logger.error(f"导出到CSV也失败: {str(csv_error)}")
    
    def load_blogger_data(self) -> pd.DataFrame:
        """
        加载博主数据文件
        
        Returns:
            pd.DataFrame: 博主数据
        """
        logger.info("开始加载博主数据...")
        
        # 获取所有博主数据Excel文件
        excel_files = []
        for file in os.listdir(self.data_dir):
            if (file.endswith(('.xlsx', '.xls', '.csv')) and 
                ('博主数据' in file or '博主' in file)):
                excel_files.append(os.path.join(self.data_dir, file))
        
        if not excel_files:
            logger.warning(f"在目录 {self.data_dir} 中未找到博主数据文件")
            return pd.DataFrame()
        
        logger.info(f"找到 {len(excel_files)} 个博主数据文件")
        
        # 合并所有文件
        dfs = []
        for file in excel_files:
            try:
                if file.endswith('.csv'):
                    df = pd.read_csv(file, encoding='utf-8')
                else:
                    df = pd.read_excel(file)
                dfs.append(df)
                logger.info(f"成功读取文件: {os.path.basename(file)}, 行数: {len(df)}")
            except Exception as e:
                logger.error(f"读取文件 {os.path.basename(file)} 失败: {str(e)}")
        
        if not dfs:
            logger.error("无法读取任何博主数据文件")
            return pd.DataFrame()
        
        # 合并DataFrame
        merged_df = pd.concat(dfs, ignore_index=True)
        logger.info(f"博主数据合并完成，总行数: {len(merged_df)}")
        
        return merged_df